import torch
from torch.nn import L1Loss, Sequential, Conv2d, MaxPool2d, Flatten, Linear, CrossEntropyLoss
from torch import nn
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10
from torchvision.transforms import ToTensor

inputs = torch.tensor([1, 2, 3], dtype=torch.float32)
targets = torch.tensor([1, 2, 5], dtype=torch.float32)

inputs = torch.reshape(inputs, (1, 1, 1, 3))
targets = torch.reshape(targets, (1, 1, 1, 3))

loss = L1Loss(reduction='sum')
result = loss(inputs, targets)

loss_mse = nn.MSELoss()
result_mse = loss_mse(inputs, targets)

print(result)
print(result_mse)

x = torch.tensor([0.1, 0.2, 0.3])
y = torch.tensor([1])
x = torch.reshape(x, (1, 3))
loss_cross = nn.CrossEntropyLoss()
result_cross = loss_cross(x, y)
print(result_cross)

dataset = CIFAR10(root="datasets", train=False, transform=ToTensor(), download=True)
dataloader = DataLoader(dataset=dataset, batch_size=64, drop_last=False)

class AModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.model = Sequential(
            Conv2d(3, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            Linear(64, 10)
        )

    def forward(self, input):
        return self.model(input)

module = AModule()
loss = CrossEntropyLoss()
optim = torch.optim.SGD(module.parameters(), 0.01)
for i in range(10):
    total_loss = 0.0
    for imgs, targets in dataloader:
        output = module(imgs)
        loss_result = loss(output, targets)
        total_loss += loss_result
        optim.zero_grad()
        loss_result.backward()
        optim.step()
    print(total_loss)
